#!/usr/bin/env python3
"""
增强安全中间件
提供全面的安全防护功能
"""

import time
import hashlib
import re
import json
import hmac
import secrets
import ipaddress
from functools import wraps
from flask import request, jsonify, g, current_app
from datetime import datetime, timedelta
import redis
from typing import Dict, Any, Optional, List, Tuple
import logging
from collections import defaultdict
import threading

class EnhancedSecurityMiddleware:
    """增强安全中间件类"""
    
    def __init__(self, redis_client):
        self.redis_client = redis_client
        self.lock = threading.Lock()
        
        # 增强的频率限制配置
        self.rate_limits = {
            'login': {'requests': 5, 'window': 300, 'block_duration': 900},      # 5次/5分钟，封禁15分钟
            'sms': {'requests': 3, 'window': 300, 'block_duration': 600},        # 3次/5分钟，封禁10分钟
            'register': {'requests': 3, 'window': 3600, 'block_duration': 1800}, # 3次/1小时，封禁30分钟
            'password_reset': {'requests': 2, 'window': 3600, 'block_duration': 3600}, # 2次/1小时，封禁1小时
            'api_call': {'requests': 1000, 'window': 3600, 'block_duration': 300}, # 1000次/1小时，封禁5分钟
            'default': {'requests': 100, 'window': 3600, 'block_duration': 300}  # 100次/1小时，封禁5分钟
        }
        
        # IP白名单和黑名单
        self.ip_whitelist = set()
        self.ip_blacklist = set()
        self.load_ip_lists()
        
        # 地理位置限制
        self.geo_restrictions = {
            'enabled': False,
            'allowed_countries': ['CN', 'US', 'GB'],  # 允许的国家代码
            'blocked_countries': ['XX']  # 禁止的国家代码
        }
        
        # 设备指纹追踪
        self.device_tracking = {
            'enabled': True,
            'max_devices_per_user': 5,
            'suspicious_threshold': 3
        }
        
        # 增强的敏感字段列表
        self.sensitive_fields = [
            'password', 'token', 'secret', 'key', 'auth', 'credential', 
            'private', 'confidential', 'api_key', 'access_token', 'refresh_token',
            'session_id', 'csrf_token', 'signature', 'hash', 'salt'
        ]
        
        # 增强的SQL注入检测模式
        self.sql_injection_patterns = [
            r"(\b(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|EXEC|UNION|TRUNCATE|REPLACE)\b)",
            r"(\b(OR|AND)\s+\d+\s*[=<>!]+\s*\d+)",
            r"(\b(OR|AND)\s+['\"]?\w+['\"]?\s*[=<>!]+\s*['\"]?\w+['\"]?)",
            r"(--|#|/\*|\*/|;)",
            r"(\bUNION\s+(ALL\s+)?SELECT\b)",
            r"(\bINTO\s+(OUT|DUMP)FILE\b)",
            r"(\bLOAD_FILE\b)",
            r"(\bCHAR\s*\(\s*\d+)",
            r"(\bCONCAT\s*\()",
            r"(\bSUBSTRING\s*\()",
            r"(\bASCII\s*\()",
            r"(\bBENCHMARK\s*\()",
            r"(\bSLEEP\s*\()",
            r"(\bWAITFOR\s+DELAY)",
            r"(\bINFORMATION_SCHEMA\b)",
            r"(\bSYS\.\w+)",
            r"(\bMSYSQL\.\w+)"
        ]
        
        # 增强的XSS检测模式
        self.xss_patterns = [
            r"<script[^>]*>.*?</script>",
            r"<iframe[^>]*>.*?</iframe>",
            r"<object[^>]*>.*?</object>",
            r"<embed[^>]*>.*?</embed>",
            r"<link[^>]*>",
            r"<meta[^>]*>",
            r"<style[^>]*>.*?</style>",
            r"javascript:",
            r"vbscript:",
            r"data:text/html",
            r"on\w+\s*=",
            r"expression\s*\(",
            r"url\s*\(",
            r"@import",
            r"<svg[^>]*>.*?</svg>",
            r"<math[^>]*>.*?</math>",
            r"<form[^>]*>.*?</form>",
            r"<input[^>]*>",
            r"<textarea[^>]*>.*?</textarea>",
            r"<select[^>]*>.*?</select>",
            r"<option[^>]*>.*?</option>",
            r"<button[^>]*>.*?</button>"
        ]
        
        # 路径遍历检测模式
        self.path_traversal_patterns = [
            r"\.\./",
            r"\.\.\\",
            r"%2e%2e%2f",
            r"%2e%2e%5c",
            r"..%2f",
            r"..%5c",
            r"%252e%252e%252f",
            r"..%252f"
        ]
        
        # 命令注入检测模式
        self.command_injection_patterns = [
            r"[;&|`$(){}[\]<>]",
            r"\b(cat|ls|pwd|id|whoami|uname|ps|netstat|ifconfig|ping|wget|curl|nc|telnet|ssh|ftp)\b",
            r"(\|\s*\w+)",
            r"(\&\&\s*\w+)",
            r"(\;\s*\w+)",
            r"(`[^`]+`)",
            r"(\$\([^)]+\))",
            r"(\$\{[^}]+\})"
        ]

    def load_ip_lists(self):
        """从Redis加载IP白名单和黑名单"""
        try:
            # 加载白名单
            whitelist_data = self.redis_client.get('security:ip_whitelist')
            if whitelist_data:
                self.ip_whitelist = set(json.loads(whitelist_data))
            
            # 加载黑名单
            blacklist_data = self.redis_client.get('security:ip_blacklist')
            if blacklist_data:
                self.ip_blacklist = set(json.loads(blacklist_data))
        except Exception as e:
            print(f"加载IP列表失败: {str(e)}")

    def save_ip_lists(self):
        """保存IP白名单和黑名单到Redis"""
        try:
            self.redis_client.setex('security:ip_whitelist', 86400, json.dumps(list(self.ip_whitelist)))
            self.redis_client.setex('security:ip_blacklist', 86400, json.dumps(list(self.ip_blacklist)))
        except Exception as e:
            print(f"保存IP列表失败: {str(e)}")

    def add_to_blacklist(self, ip: str, reason: str = "", duration: int = 3600):
        """添加IP到黑名单"""
        with self.lock:
            self.ip_blacklist.add(ip)
            self.save_ip_lists()
            
            # 记录黑名单事件
            self.log_security_event('ip_blacklisted', {
                'ip': ip,
                'reason': reason,
                'duration': duration
            })
            
            # 设置临时黑名单（带过期时间）
            self.redis_client.setex(f'security:temp_blacklist:{ip}', duration, reason)

    def is_ip_blocked(self, ip: str) -> Tuple[bool, str]:
        """检查IP是否被阻止"""
        # 检查白名单
        if ip in self.ip_whitelist:
            return False, ""
        
        # 检查永久黑名单
        if ip in self.ip_blacklist:
            return True, "IP在永久黑名单中"
        
        # 检查临时黑名单
        temp_block = self.redis_client.get(f'security:temp_blacklist:{ip}')
        if temp_block:
            return True, f"IP被临时封禁: {temp_block}"
        
        return False, ""

    def get_client_identifier(self) -> str:
        """获取客户端标识符"""
        # 优先使用真实IP
        real_ip = request.headers.get('X-Real-IP')
        if real_ip:
            return real_ip
        
        # 其次使用转发IP
        forwarded_for = request.headers.get('X-Forwarded-For')
        if forwarded_for:
            return forwarded_for.split(',')[0].strip()
        
        # 最后使用远程地址
        return request.remote_addr or 'unknown'

    def get_device_fingerprint(self) -> str:
        """生成设备指纹"""
        user_agent = request.headers.get('User-Agent', '')
        accept_language = request.headers.get('Accept-Language', '')
        accept_encoding = request.headers.get('Accept-Encoding', '')
        
        fingerprint_data = f"{user_agent}|{accept_language}|{accept_encoding}"
        return hashlib.sha256(fingerprint_data.encode()).hexdigest()[:16]

    def check_suspicious_activity(self, identifier: str) -> bool:
        """检查可疑活动"""
        # 检查短时间内的多次失败尝试
        failed_attempts_key = f"security:failed_attempts:{identifier}"
        failed_count = self.redis_client.get(failed_attempts_key)
        
        if failed_count and int(failed_count) >= 10:  # 10次失败尝试
            return True
        
        # 检查异常的请求模式
        request_pattern_key = f"security:request_pattern:{identifier}"
        pattern_data = self.redis_client.get(request_pattern_key)
        
        if pattern_data:
            pattern = json.loads(pattern_data)
            # 检查请求频率是否异常
            if pattern.get('requests_per_minute', 0) > 100:
                return True
        
        return False

    def enhanced_rate_limit_check(self, endpoint: str, identifier: str = None) -> Tuple[bool, dict]:
        """增强的频率限制检查"""
        if identifier is None:
            identifier = self.get_client_identifier()
        
        # 检查IP是否被阻止
        is_blocked, block_reason = self.is_ip_blocked(identifier)
        if is_blocked:
            return False, {
                'allowed': False,
                'reason': block_reason,
                'blocked': True
            }
        
        # 获取对应端点的限制配置
        limit_config = self.rate_limits.get(endpoint, self.rate_limits['default'])
        max_requests = limit_config['requests']
        window_seconds = limit_config['window']
        block_duration = limit_config['block_duration']
        
        # Redis键
        key = f"rate_limit:{endpoint}:{identifier}"
        
        try:
            # 获取当前计数
            current_count = self.redis_client.get(key)
            
            if current_count is None:
                # 首次请求
                self.redis_client.setex(key, window_seconds, 1)
                return True, {
                    'allowed': True,
                    'remaining': max_requests - 1,
                    'reset_time': int(time.time()) + window_seconds
                }
            else:
                current_count = int(current_count)
                
                if current_count >= max_requests:
                    # 超出限制，添加到临时黑名单
                    self.add_to_blacklist(identifier, f"频率限制超出: {endpoint}", block_duration)
                    
                    ttl = self.redis_client.ttl(key)
                    return False, {
                        'allowed': False,
                        'remaining': 0,
                        'reset_time': int(time.time()) + ttl,
                        'retry_after': ttl,
                        'blocked_until': int(time.time()) + block_duration
                    }
                else:
                    # 增加计数
                    self.redis_client.incr(key)
                    ttl = self.redis_client.ttl(key)
                    return True, {
                        'allowed': True,
                        'remaining': max_requests - current_count - 1,
                        'reset_time': int(time.time()) + ttl
                    }
                    
        except Exception as e:
            # Redis错误时允许请求通过，但记录错误
            print(f"Rate limit check failed: {str(e)}")
            return True, {'allowed': True, 'error': str(e)}

    def enhanced_input_validation(self, data: Any, max_depth: int = 5) -> Tuple[bool, str, List[str]]:
        """增强的输入验证"""
        threats_detected = []
        
        if max_depth <= 0:
            return False, "数据嵌套层级过深", ["DEEP_NESTING"]
        
        if isinstance(data, dict):
            # 检查字典大小
            if len(data) > 100:
                return False, "字典键值对数量超出限制", ["LARGE_DICT"]
            
            for key, value in data.items():
                # 检查键名
                key_safe, key_threats = self._check_string_security(str(key))
                if not key_safe:
                    threats_detected.extend(key_threats)
                    return False, f"不安全的键名: {key}", threats_detected
                
                # 递归检查值
                is_valid, message, sub_threats = self.enhanced_input_validation(value, max_depth - 1)
                threats_detected.extend(sub_threats)
                if not is_valid:
                    return False, message, threats_detected
                    
        elif isinstance(data, list):
            if len(data) > 1000:  # 限制数组长度
                return False, "数组长度超出限制", ["LARGE_ARRAY"]
            
            for item in data:
                is_valid, message, sub_threats = self.enhanced_input_validation(item, max_depth - 1)
                threats_detected.extend(sub_threats)
                if not is_valid:
                    return False, message, threats_detected
                    
        elif isinstance(data, str):
            if len(data) > 10000:  # 限制字符串长度
                return False, "字符串长度超出限制", ["LARGE_STRING"]
            
            string_safe, string_threats = self._check_string_security(data)
            threats_detected.extend(string_threats)
            if not string_safe:
                return False, "检测到不安全的字符串内容", threats_detected
                
        return True, "验证通过", threats_detected

    def _check_string_security(self, text: str) -> Tuple[bool, List[str]]:
        """检查字符串安全性"""
        threats = []
        
        # SQL注入检测
        for pattern in self.sql_injection_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                threats.append("SQL_INJECTION")
                break
        
        # XSS检测
        for pattern in self.xss_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                threats.append("XSS")
                break
        
        # 路径遍历检测
        for pattern in self.path_traversal_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                threats.append("PATH_TRAVERSAL")
                break
        
        # 命令注入检测
        for pattern in self.command_injection_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                threats.append("COMMAND_INJECTION")
                break
        
        # 检查是否包含null字节
        if '\x00' in text:
            threats.append("NULL_BYTE")
        
        # 检查是否包含控制字符
        if any(ord(c) < 32 and c not in '\t\n\r' for c in text):
            threats.append("CONTROL_CHARS")
        
        return len(threats) == 0, threats

    def enhanced_data_sanitization(self, data: Any) -> Any:
        """增强的数据清理"""
        if isinstance(data, dict):
            return {key: self.enhanced_data_sanitization(value) for key, value in data.items()}
        elif isinstance(data, list):
            return [self.enhanced_data_sanitization(item) for item in data]
        elif isinstance(data, str):
            # 移除危险字符
            sanitized = data.replace('\x00', '')  # 移除null字节
            
            # 移除控制字符（保留制表符、换行符、回车符）
            sanitized = ''.join(c for c in sanitized if ord(c) >= 32 or c in '\t\n\r')
            
            # 移除脚本标签
            sanitized = re.sub(r'<script[^>]*>.*?</script>', '', sanitized, flags=re.IGNORECASE | re.DOTALL)
            
            # 移除事件处理器
            sanitized = re.sub(r'on\w+\s*=\s*["\'][^"\']*["\']', '', sanitized, flags=re.IGNORECASE)
            
            # 移除javascript:和vbscript:
            sanitized = re.sub(r'(javascript|vbscript):', '', sanitized, flags=re.IGNORECASE)
            
            return sanitized
        else:
            return data

    def log_security_event(self, event_type: str, details: dict):
        """记录安全事件"""
        log_data = {
            'timestamp': datetime.utcnow().isoformat(),
            'event_type': event_type,
            'client_ip': self.get_client_identifier(),
            'user_agent': request.headers.get('User-Agent', ''),
            'endpoint': request.endpoint,
            'method': request.method,
            'device_fingerprint': self.get_device_fingerprint(),
            'details': details
        }
        
        # 存储到Redis
        log_key = f"security_log:{int(time.time())}:{secrets.token_hex(8)}"
        try:
            self.redis_client.setex(log_key, 86400 * 30, json.dumps(log_data))  # 保存30天
        except Exception as e:
            print(f"Failed to log security event: {str(e)}")

    def get_enhanced_security_headers(self) -> dict:
        """获取增强的安全响应头"""
        return {
            'X-Content-Type-Options': 'nosniff',
            'X-Frame-Options': 'DENY',
            'X-XSS-Protection': '1; mode=block',
            'Strict-Transport-Security': 'max-age=31536000; includeSubDomains; preload',
            'Content-Security-Policy': "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self'; connect-src 'self'; media-src 'none'; object-src 'none'; child-src 'none'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'",
            'Referrer-Policy': 'strict-origin-when-cross-origin',
            'Permissions-Policy': 'geolocation=(), microphone=(), camera=(), payment=(), usb=(), magnetometer=(), gyroscope=(), speaker=(), vibrate=(), fullscreen=()',
            'X-Permitted-Cross-Domain-Policies': 'none',
            'X-Download-Options': 'noopen',
            'Cache-Control': 'no-cache, no-store, must-revalidate',
            'Pragma': 'no-cache',
            'Expires': '0',
            'X-Robots-Tag': 'noindex, nofollow, nosnippet, noarchive'
        }

# 增强的装饰器
def enhanced_rate_limit(endpoint: str):
    """增强的请求频率限制装饰器"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not hasattr(g, 'enhanced_security_middleware'):
                return jsonify({
                    'success': False,
                    'message': '增强安全中间件未初始化',
                    'timestamp': datetime.utcnow().isoformat()
                }), 500
            
            # 检查频率限制
            allowed, info = g.enhanced_security_middleware.enhanced_rate_limit_check(endpoint)
            
            if not allowed:
                # 记录频率限制事件
                g.enhanced_security_middleware.log_security_event('enhanced_rate_limit_exceeded', {
                    'endpoint': endpoint,
                    'limit_info': info
                })
                
                response = jsonify({
                    'success': False,
                    'message': '请求过于频繁，已被临时封禁',
                    'retry_after': info.get('retry_after', 60),
                    'blocked_until': info.get('blocked_until'),
                    'timestamp': datetime.utcnow().isoformat()
                })
                response.status_code = 429
                response.headers['Retry-After'] = str(info.get('retry_after', 60))
                return response
            
            # 添加频率限制信息到响应头
            response = f(*args, **kwargs)
            if hasattr(response, 'headers'):
                response.headers['X-RateLimit-Remaining'] = str(info.get('remaining', 0))
                response.headers['X-RateLimit-Reset'] = str(info.get('reset_time', 0))
            
            return response
        return decorated_function
    return decorator

def enhanced_validate_input():
    """增强的输入验证装饰器"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not hasattr(g, 'enhanced_security_middleware'):
                return jsonify({
                    'success': False,
                    'message': '增强安全中间件未初始化',
                    'timestamp': datetime.utcnow().isoformat()
                }), 500
            
            # 验证请求数据
            if hasattr(g, 'decrypted_data') and g.decrypted_data:
                is_valid, message, threats = g.enhanced_security_middleware.enhanced_input_validation(g.decrypted_data)
                
                if not is_valid:
                    # 记录安全事件
                    g.enhanced_security_middleware.log_security_event('enhanced_invalid_input', {
                        'validation_error': message,
                        'threats_detected': threats,
                        'data_sample': str(g.decrypted_data)[:200]  # 只记录前200字符
                    })
                    
                    # 如果检测到严重威胁，添加到黑名单
                    serious_threats = ['SQL_INJECTION', 'COMMAND_INJECTION', 'PATH_TRAVERSAL']
                    if any(threat in serious_threats for threat in threats):
                        client_ip = g.enhanced_security_middleware.get_client_identifier()
                        g.enhanced_security_middleware.add_to_blacklist(
                            client_ip, 
                            f"严重安全威胁: {', '.join(threats)}", 
                            3600  # 封禁1小时
                        )
                    
                    return jsonify({
                        'success': False,
                        'message': '输入数据验证失败',
                        'timestamp': datetime.utcnow().isoformat()
                    }), 400
                
                # 清理数据
                g.decrypted_data = g.enhanced_security_middleware.enhanced_data_sanitization(g.decrypted_data)
            
            return f(*args, **kwargs)
        return decorated_function
    return decorator